# -*- coding: utf-8 -*-
# Author: Sun Jiawei
# E-mail: sunjiawei@tbea.com
#
import os
import warnings
warnings.filterwarnings('ignore')

import pywt
import pandas as pd
import matplotlib.pyplot as plt

from tqdm import tqdm
from sklearn.metrics import mean_absolute_percentage_error as mape
from statsmodels.tsa.stattools import adfuller
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
from statsmodels.tsa.arima.model import ARIMA

# 原始数据处理
excel_dir = '/Volumes/工作/售电交易一体化/图迹平台'
print(os.listdir(excel_dir))

file_names = {
    '2021年07月01日-2021年07月31日等31天价格预测(96点).xlsx': pd.date_range('7/1/2021', '7/31/2021'),
    '2021年08月01日-2021年08月31日等31天价格预测(96点).xlsx': pd.date_range('8/1/2021', '8/31/2021'),
    '2021年09月01日-2021年09月30日等30天价格预测(96点).xlsx': pd.date_range('9/1/2021', '9/30/2021'),
    '2021年10月01日-2021年10月31日等31天价格预测(96点).xlsx': pd.date_range('10/1/2021', '10/31/2021'),
    '2021年11月01日-2021年11月30日等30天价格预测(96点).xlsx': pd.date_range('11/1/2021', '11/30/2021')
}
columns = ['timestamp', 'tg实时预测', 'tg日前预测', '实时实际', '日前实际']
df = pd.DataFrame(columns=columns)

for i in tqdm(file_names):
    sheet_names = file_names[i]
    file_name = os.path.join(excel_dir, i)
    for sheet_name in tqdm(sheet_names):
        sheet_name = sheet_name.strftime('%Y-%m-%d')
        timestamp = pd.date_range(start=sheet_name, periods=96, freq='15min')
        _df = pd.read_excel(file_name, sheet_name=sheet_name)
        _df.columns = columns
        _df.timestamp = timestamp
        df = df.append(_df, ignore_index=True)
df.to_csv('./data/图迹数据-7月至11月.csv', index=False)

# 建模并应用


class GeekDS:
    def __init__(self, csv_path: str) -> None:
        """Initialize a `DataSet` object."""
        self.data = pd.read_csv(csv_path)
        self.data['timestamp'] = pd.to_datetime(self.data['timestamp'])
        self.data = self.data.set_index('timestamp')
        self.data['hour'] = self.data.index.hour
        self.data['minute'] = self.data.index.minute
        self.data['week'] = self.data.index.weekday + 1

    def get_date_range_data(self, start_year: str, start_month: str, start_day: str,
                            end_year: str, end_month: str, end_day: str) -> pd.DataFrame:
        """获取指定时间短内的数据.

        :param start_year: 开始年
        :param start_month: 开始月
        :param start_day: 开始日
        :param end_year: 终止年
        :param end_month: 终止月
        :param end_day: 终止日
        :return:
        """
        start = pd.Timestamp(start_year + '-' + start_month + '-' + start_day)
        end = pd.Timestamp(end_year + '-' + end_month + '-' + end_day)
        return self.data.truncate(before=start, after=end)

    def get_panel_data(self, start_year: str, start_month: str, end_year: str, end_month: str,
                       hour: str, minute: str) -> pd.DataFrame:
        """在一个指定的时间段内,获取每一天指定时刻的数据.

        由于每个月的天数不一样,因此,以下个月的1号为判定日期来确定时间段的终止年月,如有必要,再对数据进行删减处理.

        :param start_year: 时间段的起始年
        :param start_month: 时间段的起始月
        :param end_year: 时间段的终止年
        :param end_month: 时间段的终止月
        :param hour: 指定时刻的小时
        :param minute: 指定时刻的分钟
        :return:
        """
        start = pd.Timestamp(start_year + '-' + start_month + '-1')
        end = pd.Timestamp(end_year + '-' + end_month + '-1')
        df = self.data.truncate(before=start, after=end)
        return df[(df.hour == int(hour)) & (df.minute == int(minute))]


# 数据可视化部分
ds = GeekDS('./data/图迹数据-7月至11月.csv')

h1, m1 = '1', '0'
h2, m2 = '20', '0'
df1 = ds.get_panel_data('2021', '7', '2021', '12', h1, m1)
df2 = ds.get_panel_data('2021', '7', '2021', '12', h2, m2)
x = df1.index
y1 = df1['日前实际']
y2 = df2['日前实际']

fig, ax = plt.subplots()
ax.set_title('价格曲线')
ax.set_ylabel('价格(元/MWh)')
ax.set_xlabel('日期')
ax.plot(x, y1, label=f'{h1}:{m1}价格')
ax.plot(x, y2, label=f'{h2}:{m2}价格')
ax.legend()
plt.show()

# 用7-9月的数据来建模,10月的数据来评估模型的准确性
df = ds.get_panel_data('2021', '7', '2021', '11', h1, m1)
df_train = df.iloc[:-31, :]
df_val = df.iloc[-31:, :]
wavelet_name = 'db1'
w = pywt.Wavelet(wavelet_name)
a1, _ = pywt.dwt(df_train['日前实际'].values, w, mode=pywt.Modes.smooth)
a1 = pywt.waverec([a1, None], w)
a1 = pd.DataFrame(a1, index=pd.date_range('7/1/2021', '9/30/2021'))

fig = plt.figure()
ax_main = fig.add_subplot(3, 1, 1)
ax_main.set_title('原始电价序列')
ax_main.plot(df_train['日前实际'].values)
ax_main.set_xlim(0, len(df_train['日前实际'].values)-1)

ax = fig.add_subplot(3, 1, 2)
ax.plot(a1.values, 'r')
ax.set_xlim(0, len(a1) - 1)
ax.set_ylabel("概貌序列")
plt.show()

diff_a = a1.diff().iloc[1:]
# 可以对差分后的序列进行adf检验,以判断是否平稳
# adf检验的原假设是序列为单位根,备择假设为序列是平稳的
adf_res = adfuller(diff_a)
print(f"检验的p值为:{adf_res[1]}")  # noqa
if adf_res[1] < 0.05:  # noqa
    print("序列是平稳的")
else:
    print("序列是非平稳的!")

# 目测arma模型的阶数
plot_acf(diff_a)
plot_pacf(diff_a, method='ywm')
plt.show()

model = ARIMA(diff_a, order=(1, 0, 3)).fit()

preds = model.predict(pd.to_datetime('10/1/2021'), pd.to_datetime('10/31/2021'), dynamic=True)
preds.iloc[0] += df_train['日前实际'].values[-1]
preds = preds.cumsum().values
y_true = df_val['日前实际'].values
# 对日前实际值为0的进行处理
y_true_ = []
y_pred_ = []
y_pred_tg = []
for i in range(len(y_true)):
    if y_true[i] != 0:
        y_true_.append(y_true[i])
        y_pred_.append(preds[i])
        y_pred_tg.append(df_val['tg日前预测'].values[i])
    else:
        continue

print(mape(y_true_, y_pred_))
print(mape(y_true_, y_pred_tg))

plt.plot(y_true_)
plt.plot(y_pred_)
plt.plot(y_pred_tg)
plt.show()
